///
/// Copyright (C) 2014-2015, Cyberhaven
/// All rights reserved.
///
/// Licensed under the Cyberhaven Research License Agreement.
///
#include <s2e/cpu.h>
#include <s2e/ConfigFile.h>
#include <s2e/S2E.h>
#include <s2e/Utils.h>
#include <capstone/capstone.h>

#include "ForkLimiter.h"

namespace s2e {
namespace plugins {

S2E_DEFINE_PLUGIN(ForkLimiter, "Limits how many times each instruction in a module can fork", "");

void ForkLimiter::initialize() {
    m_detector = s2e()->getPlugin<ModuleExecutionDetector>();

    s2e()->getCorePlugin()->onTimer.connect(sigc::mem_fun(*this, &ForkLimiter::onTimer));

    s2e()->getCorePlugin()->onProcessForkDecide.connect(sigc::mem_fun(*this, &ForkLimiter::onProcessForkDecide));

    // Limit of forks per program counter, -1 means don't care
    bool ok;
    m_limit = s2e()->getConfig()->getInt(getConfigKey() + ".maxForkCount", 10, &ok);

    m_underConstraint = s2e()->getConfig()->getBool(getConfigKey() + ".underConstraint");
    m_fpsMode = s2e()->getConfig()->getBool(getConfigKey() + ".fpsMode");
    m_generationControl = s2e()->getConfig()->getInt(getConfigKey() + ".generationControl", 10000);
    ConfigFile *cfg = s2e()->getConfig();
	ConfigFile::string_list pollingEntries = cfg->getListKeys(
			getConfigKey() + ".forkRanges");
	if (pollingEntries.size() == 0) {
		Range rg;
		rg.start = 0x00000000;
		rg.end = 0x00000000;
		m_forkRanges.insert(rg);
	}
	foreach2(it, pollingEntries.begin(), pollingEntries.end())
	{
		std::stringstream ss1;
		ss1 << getConfigKey() << ".forkRanges" << "." << *it;
		ConfigFile::integer_list il = cfg->getIntegerList(ss1.str());
		if (il.size() != 2) {
			s2e()->getDebugStream() << "Range entry " << ss1.str()
					<< " must be of the form {startPc, endPc} format" << '\n';
			continue;
		}

		bool ok = false;
		uint64_t start = cfg->getInt(ss1.str() + "[1]", 0, &ok);
		if (!ok) {
			s2e()->getDebugStream() << "ForkLimiter could not read "
					<< ss1.str() << "[0]" << '\n';
			continue;
		}

		uint64_t end = cfg->getInt(ss1.str() + "[2]", 0, &ok);
		if (!ok) {
			s2e()->getDebugStream() << "ForkLimiter could not read "
					<< ss1.str() << "[1]" << '\n';
			continue;
		}
		Range rg;
		rg.start = start;
		rg.end = end;
		m_forkRanges.insert(rg);
	}

    if ((int) m_limit != -1) {
        if (m_detector) {
            s2e()->getCorePlugin()->onStateForkDecide.connect(sigc::mem_fun(*this, &ForkLimiter::onStateForkDecide));

            s2e()->getCorePlugin()->onStateFork.connect(sigc::mem_fun(*this, &ForkLimiter::onFork));
        } else if (ok) {
            getDebugStream() << "maxForkCount requires ModuleExecutionDetector\n";
            exit(-1);
        }
    }

    // Wait 5 seconds before allowing an S2E instance to fork
    m_processForkDelay = s2e()->getConfig()->getInt(getConfigKey() + ".processForkDelay", 5);

    m_timerTicks = 0;
}

void ForkLimiter::onStateForkDecide(S2EExecutionState *state, bool *doFork) {
    auto module = m_detector->getCurrentDescriptor(state);
	auto moduleall = m_detector->getAllDescriptor(state);
	state->m_fpsMode = m_fpsMode;
	if (moduleall) {
		size_t pos = moduleall->Name.find(std::string("s2ecmd"));
		if (pos != std::string::npos) {
			return;
		}
	}

    if (!module) {
    	if(m_underConstraint){
    		state->underConstraint = true;
    		*doFork = false;
    	}
    	if(m_fpsMode)
    	{
    		*doFork = false;
    	}
        return;
    }

    uint64_t curPc;
    if (!module->ToNativeBase(state->regs()->getPc(), curPc)) {
    	getDebugStream(state) << "Could not get relative pc for module " << module->Name << "\n";
    	if(m_underConstraint){
    		state->underConstraint = true;
    		*doFork = false;
    	}
    	if(m_fpsMode)
    	{
    		*doFork = false;
    	}
        return;
    }

    TranslationBlock *tb = state->getTb();
    if(m_fpsMode){
    	bool allowFork = false;
    	if (tb){
    		uint64_t bb_start_pc = ((uint64_t) -1);
    		if(module->ToNativeBaseRebased(tb->pc, bb_start_pc)){
				foreach2(rgit, m_forkRanges.begin(), m_forkRanges.end())
				{
					if (bb_start_pc == (*rgit).start) {
						allowFork = true;
						break;
					}
				}
    		}
    		s2e()->getInfoStream(state) << "Forkable at fps_bb_start_pc=" << hexval(bb_start_pc)<< "\n";
    	}
    	if(allowFork){
    		getDebugStream(state) << "Allow StateForkDecide Native pc= " << hexval(curPc)<< "\n";
    	}else{
    		getDebugStream(state) << "Disable StateForkDecide Native pc= " << hexval(curPc)<< "\n";
        	if(m_underConstraint){//??
        		state->underConstraint = true;
        	}
    		*doFork = false;
    	}
    }
	int flags = 0; // 32-bit code by default
	if (tb) {
		switch (state->getPointerSize()) {
		case 4:
			flags = 0;
			break;
		case 8:
			flags = 2;
			break;
		default:
			flags = 0;//pabort("Not supported code");
		}
	}

	csh handle;
	cs_mode mode;
	switch (flags) {
	case 0:
		mode = CS_MODE_32;
		break;
	case 1:
		mode = CS_MODE_16;
		break;
	case 2:
		mode = CS_MODE_64;
		break;
	default:
		mode = CS_MODE_32;
		break;
	}

	if (cs_open(CS_ARCH_X86, mode, &handle) == CS_ERR_OK) {
		cs_insn *insn;
		uint64_t pcaddress = state->regs()->getPc();
		size_t count;
		unsigned char code[15];
		size_t codesize = 15;
		if (state->mem()->read(pcaddress, code, codesize)) {
			count = cs_disasm(handle, code, codesize, pcaddress, 0, &insn);
			if (count) {
				//printf("Disasm:");
				//printf("0x%08lx:\t%s\t\t%s\n", insn[0].address,
				//		insn[0].mnemonic, insn[0].op_str);
				if(insn[0].mnemonic[0]=='f'){
					getDebugStream(state) << "Fork disabled for instruction " << insn[0].mnemonic << " for module " << module->Name  << " at "<< hexval(curPc) << "\n";
					if(m_underConstraint){//??
						state->underConstraint = true;
					}
					*doFork = false;
				}
				cs_free(insn, count);
			} else {
			}
		}
		cs_close(&handle);
	}

    if (m_forkCount[module->Name][curPc] > m_limit) {
        *doFork = false;
    }
}

void ForkLimiter::onFork(S2EExecutionState *state, const std::vector<S2EExecutionState *> &newStates,
                         const std::vector<klee::ref<klee::Expr>> &newConditions) {
//	getDebugStream(state) << "ForkLimiter::onFork id=" << state->getID() << " " << state->m_generationCount << "\n";
	if(state && state->m_generationCount > m_generationControl){ state->forkDisabled = true; getDebugStream(state) << "ForkDisabled for id=" << state->getID() << "\n";}
	for (unsigned i = 0; i < newStates.size(); i++) {
		if(newStates[i] && newStates[i]->m_generationCount > m_generationControl){
			newStates[i]->forkDisabled = true;
			s2e()->getExecutor()->terminateState(*newStates[i],
					"terminate write generation case.");
		}
	}

	auto module = m_detector->getCurrentDescriptor(state);
    if (!module) {
        return;
    }

    uint64_t curPc;
    if (!module->ToNativeBase(state->regs()->getPc(), curPc)) {
        getDebugStream(state) << "Could not get relative pc for module " << module->Name << "\n";
        return;
    }else{
    	getDebugStream(state) << "Forking at relative pc for module " << module->Name  << " at "<< hexval(curPc) << "\n";
    }
    ++m_forkCount[module->Name][curPc];
}

void ForkLimiter::onProcessForkDecide(bool *proceed) {
    // Rate-limit forking
    if (m_timerTicks < m_processForkDelay) {
        *proceed = false;
        return;
    }

    m_timerTicks = 0;
}

void ForkLimiter::onTimer() {
    ++m_timerTicks;
}

} // namespace plugins
} // namespace s2e
